Greptile Summary

This PR adds a NVTE_JAX_ENFORCE_V2_GROUPED_GEMM environment variable that allows users to assert the V2 (CUDA-graphable) grouped GEMM path is taken, turning silent fallbacks into loud RuntimeErrors useful for debugging and CI enforcement.

Key changes in transformer_engine/jax/cpp_extensions/gemm.py:

• Captures _v2_grouped_gemm_available_reason at import time so the cuBLAS version error can be surfaced in the enforcement message.
• Adds a @cache-decorated _should_enforce_v2_grouped_gemm() helper that reads the env var exactly once per process.
• Extends _can_use_v2_grouped_gemm with three guard-point checks: cuBLAS version availability, SM compute capability, and dtype/quantization/bias restrictions.
• One minor style note: get_device_compute_capability(0) is called twice in the compute-capability guard (once for the comparison and once inside the error message string).

Confidence Score: 5/5

• Safe to merge — no logic or runtime errors introduced; the feature is additive and opt-in via an env var.
• All three enforcement guard points are logically correct and consistent with the existing fallback logic. Error messages are clear and include diagnostic context. The only finding is a trivial style suggestion (double get_device_compute_capability(0) call) that does not affect correctness or reliability.
• No files require special attention.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[_can_use_v2_grouped_gemm called] --> B[read enforce_v2_gmm via cached env-var]
    B --> C{_v2_grouped_gemm_available?}
    C -- No --> D{enforce_v2_gmm?}
    D -- Yes --> E[raise RuntimeError with reason string]
    D -- No --> F[return False]
    C -- Yes --> G{device SM < 100?}
    G -- Yes --> H{enforce_v2_gmm?}
    H -- Yes --> I[raise RuntimeError with compute cap]
    H -- No --> J[return False]
    G -- No --> K{NO_SCALING + BF16 + no bias?}
    K -- Yes --> L[return True  → use V2 path]
    K -- No --> M{enforce_v2_gmm?}
    M -- Yes --> N[raise RuntimeError with dtype/bias/mode]
    M -- No --> O[return False  → use V1 path]

_{Reviews (1): Last reviewed commit: "TE GMM v2 enforcement env var" | Re-trigger Greptile}

/te-ci

Closing as this change merged as part of #2749